import numpy as np
import pandas as pd
import torch
from torch.utils.data import TensorDataset
from nlpx.model.classifier import TextCNNClassifier
from nlpx.text_token import PaddingTokenizer
from nlpx.training import Trainer
from nlpx.text_token.utils import get_df_text_labels
from nlpx.model.wrapper import ModelWrapper, FastModelWrapper


if __name__ == '__main__':
	df = pd.read_csv('~/project/python/parttime/text_gcn/data/北方地区不安全事件统计20240331.csv', encoding='GBK')
	texts, labels, classes = get_df_text_labels(df, text_col='故障描述', label_col='故障标志')
	
	tokenizer = PaddingTokenizer(texts=texts, cut_type='char')
	ids = tokenizer.batch_encode(texts)
	ids = torch.LongTensor(ids)
	dataset = TensorDataset(ids, torch.tensor(labels, dtype=torch.long))
	
	model = TextCNNClassifier(100, vocab_size=tokenizer.vocab_size, num_classes=len(classes))
	
	# trainer = Trainer(epochs=500)
	# trainer.train(model, dataset)
	
	model_wrapper = ModelWrapper(model)
	model_wrapper.train(dataset, early_stopping_rounds=2, show_progress=True)
	
	# model_wrapper = FastModelWrapper()
	# model_wrapper.train(model, ids, labels)

	test_texts = ['昆明航B737飞机执行昆明至西安航班 西安机场进近过程中  机组将飞机襟翼位置设置错误触发近地警告事件',
				  '海航 B737飞机执行汉中至广州航班 飞机起飞后发现汉中过站期间货舱有730KG货物未卸机']
	# inputs = tokenize_vec.encode_plus(test_texts)
	# print(model.predict_classes_proba(inputs))
	
	test_ids = tokenizer.batch_encode(test_texts)
	test_ids = torch.LongTensor(test_ids)
	print(model_wrapper.logits(test_ids))
	

	